import Yices
import Data.List

data Semester = None | First | Second deriving (Bounded, Enum, Show)
instance Named Semester where sname _ = "semester"
instance Refl Semester where refl = stdRefl
instance Scalar Semester

-- Some constants
courses = [1..10]

-- Course -> Semester mapping
s' :: Ident (Nat -> Semester)
(s',s) = ident1 "s"

-- ECTs per course
ects :: Nat -> Nat
ects = (!!) [0,4,5,4,5,4,6,6,5,6,6] . fromIntegral

-- The main script
main = printScript
    -- Declare the semester type and variables
    [ deftype (undefined :: Semester)
    , define s'
    , assert $ atLeastTwoFrom [1..4] /\ atLeastTwoFrom [5..7] /\ atLeastTwoFrom [8..10]
    , assert $ prereq 8 1 /\ prereq 9 2 /\ prereq 9 3
    , assert $ incompatible 7 10
    , assert $ semesterECTs First == nat 30 /\ semesterECTs Second == nat 10
    ]
  where
    -- DSL for the requirements
    atLeastTwoFrom cs  = disj [non (i `during` None) /\ non (j `during` None) | i <- cs, j <- cs, j > i]
    semesterECTs   s   = sum [pseudo (c `during` s) (ects c) | c <- courses]
    prereq         c p = (c `during` None) \/ ((c `during` Second) /\ (p `during` First))
    incompatible   c d = (c `during` None) \/ (d `during` None)
    c `during` sem     = eq' ! s c ! lit sem

    -- General-purpose Yices operators
    pseudo x n  = ifte' ! x ! nat n ! nat 0
    sum = foldr (+) (nat 0)
    x + y = plus' ! x ! y
    x == y = Yices.eq' ! x ! y

    -- Specialized Yices identifiers
    ifte' :: Ident (Bool -> Nat -> Nat -> Nat)
    ifte' = ident "if"
    eq' :: Ident (Semester -> Semester -> Bool)
    eq' = ident "="
